# faster-rcnn_r50_fpn_1x_coco.py 继承配置
# _base_ = [
#     '../_base_/models/faster-rcnn_r50_fpn.py',
#     '../_base_/datasets/coco_detection.py',
#     '../_base_/schedules/schedule_1x.py', '../_base_/default_runtime.py'
# ]
import os

# 继承配置(_base_对象只能为列表或者字符串)
_base_ = 'faster-rcnn_r50_fpn_1x_coco.py'  # 继承同目录下的 faster-rcnn_r50_fpn_1x_coco.py 配置

# 重写训练轮次
train_cfg = dict(type='EpochBasedTrainLoop', max_epochs=150, val_interval=1)

# learning rate
param_scheduler = [
    dict(
        type='LinearLR', start_factor=0.001, by_epoch=False, begin=0, end=500),
    dict(
        type='MultiStepLR',
        begin=0,
        end=12,
        by_epoch=True,
        milestones=[80, 110],  # 在 80，110轮次改变学习率
        gamma=0.1)
]

# 我们还需要更改 head 中的 num_classes 以匹配数据集中的类别数
model = dict(
    roi_head=dict(
        type='StandardRoIHead',
        bbox_roi_extractor=dict(
            type='SingleRoIExtractor',
            roi_layer=dict(type='RoIAlign', output_size=7, sampling_ratio=0),
            out_channels=256,
            featmap_strides=[4, 8, 16, 32]),
        bbox_head=dict(
            type='Shared2FCBBoxHead',
            in_channels=256,
            fc_out_channels=1024,
            roi_feat_size=7,
            num_classes=2,
            bbox_coder=dict(
                type='DeltaXYWHBBoxCoder',
                target_means=[0., 0., 0., 0.],
                target_stds=[0.1, 0.1, 0.2, 0.2]),
            reg_class_agnostic=False,
            loss_cls=dict(
                type='CrossEntropyLoss', use_sigmoid=False, loss_weight=1.0),
            loss_bbox=dict(type='L1Loss', loss_weight=1.0)))
)

# 修改数据集相关配置（绝对路径）
data_root = '/root/mmdetection/dataset/SDB_9K_COCO/'
metainfo = {
    'classes': ('drone', 'bird'),
    'palette': [
        (253, 58, 52), (253, 159, 148),
    ]
}
train_dataloader = dict(
    batch_size=4,
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='train/annotations/SDB_9K_train.json',
        data_prefix=dict(img='train/images')))
val_dataloader = dict(
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='val/annotations/SDB_9K_val.json',
        data_prefix=dict(img='val/images')))
test_dataloader = dict(
    dataset=dict(
        data_root=data_root,
        metainfo=metainfo,
        ann_file='test/annotations/SDB_9K_test.json',
        data_prefix=dict(img='test/images')))

# 修改评价指标相关配置
val_evaluator = dict(ann_file=os.path.join(data_root, 'val/annotations/SDB_9K_val.json'))
test_evaluator = dict(ann_file=os.path.join(data_root, 'test/annotations/SDB_9K_test.json'))

# 使用预训练的 Faster R-CNN 模型权重来做初始化，可以提高模型性能
load_from = 'https://download.openxlab.org.cn/models/mmdetection/FasterR-CNN/weight/faster-rcnn_r50_fpn_1x_coco'
